#!/usr/bin/env python
"""
CTPRPCPickle

serialization/deserialization used by CTP-RPC v2 protocol
"""

import os, sys
import struct
from array import array
from datetime import datetime
from cPickle import loads, dumps
from StringIO import StringIO


_little_endian = ( sys.byteorder == 'little' )

class SerializationError ( Exception ): pass

# *******************************************************
# CTPTempFile class
class CTPTempFile ( object ):
	def __init__ ( self, name ):
		super ( CTPTempFile, self ).__init__()
		self._file = os.tmpfile()
		self.name = name


	# *******************************************************
	# __del__() - distruttore, chiude il file
	def __del__ ( self ):
		self.close()

	
	# *******************************************************
	# __getattr__() - restituisce gli attributi di file
	def __getattr__ ( self, name ):
		f = self.__dict__.get ( "_file" )
		if not f: raise AttributeError
		return getattr ( f, name )


class CTPRPCSerializer ( object ):
	def __init__ ( self, relax_serialize = 0 ):
		super ( CTPRPCSerializer, self ).__init__ ()

		self._relaxed = relax_serialize


	def _serialize_string ( self, s ):
		ret = 's'
		ret += struct.pack ( ">I", len ( s ) )
		return ret + s


	def _serialize_unicode ( self, u ):
		ret = 'S'

		a = array ( 'u', u )
		if _little_endian: a.byteswap ()
		s = a.tostring ()
		
		ret += struct.pack ( ">I", len ( s ) )
		return ret + s


	def _serialize_int8 ( self, i, unsigned = 0 ):
		ret = ( unsigned and 'B' or 'b' )
		ret += struct.pack ( ret, i )
		return ret


	def _serialize_int16 ( self, i, unsigned = 0 ):
		ret = ( unsigned and 'H' or 'h' )
		ret += struct.pack ( ">" + ret, i )
		return ret


	def _serialize_int32 ( self, i, unsigned = 0 ):
		ret = ( unsigned and 'I' or 'i' )
		ret += struct.pack ( ">" + ret, i )
		return ret


	def _serialize_int64 ( self, i, unsigned = 0 ):
		ret = ( unsigned and 'L' or 'l' )
		ret += struct.pack ( ">" + ( unsigned and 'Q' or 'q' ), i )
		return ret


	def _serialize_bigint ( self, i ):
		ret = StringIO ()
		ret.write ( (i < 0) and '-' or '+' )
		h = hex ( i ).lower ()
		if h [ -1 ]  == 'l': h = h [ : -1 ]
		if h [ : 2 ] == '0x': h = h [ 2 : ]
		elif h [ : 3 ] == '-0x': h = h [ 3 : ]

		l = len ( h )
		if l % 2:
			h = '0' + h
			l += 1

		ret.write ( struct.pack ( '>I', l / 2 ) )

		for i in xrange ( 0, l, 2 ):
			byte = h [ i : i + 2 ]
			ret.write ( struct.pack ( 'B', int ( byte, 16 ) ) )
		
		return ret.getvalue ()


	def _serialize_bool ( self, b ):
		return b and '1' or '0'


	def _serialize_datetime ( self, d ):
		ret = "d"
		ret += struct.pack ( ">HBBBBBH", d.year, d.month, d.day, d.hour, d.minute, d.second, d.microsecond / 1000 )
		return ret


	def _serialize_float ( self, f, single ):
		ret = ( single and 'f' or 'F' )
		ret += struct.pack ( ">" + ( single and 'f' or 'd' ), f )
		return ret


	def _serialize_list ( self, l ):
		ret = StringIO ()
		ret.write ( "[" )
		ret.write ( struct.pack ( ">I", len ( l ) ) )
		for item in l:
			ret.write ( self._serialize ( item ) )
		return ret.getvalue ()


	def _serialize_dict ( self, d ):
		ret = StringIO ()
		ret.write ( "{" )
		ret.write ( struct.pack ( ">I", len ( d ) ) )
		for k, v in d.iteritems ():
			ret.write ( self._serialize ( k ) )
			ret.write ( self._serialize ( v ) )
		return ret.getvalue ()


	def _serialize_file ( self, f ):
		rname = self._serialize_string ( f.name )
		ret = '/' + rname [ 1 : ]
		self.files.append ( f )
		return ret


	def _serialize_python ( self, obj ):
		ser = dumps ( obj )
		ret = "X"
		ret += struct.pack ( ">I", len ( ser ) )
		ret += ser
		return ret
		

	def serialize ( self, obj, flags = None ):
		"""
		serialize()
			flags:	'u' -> force unsigned
				'b' -> serialize as int8 (requires int object)
				'h' -> serialize as int16 (requires int object)
				'i' -> serialize as int32 (requires long object)
				'l' -> serialize as int64 (requires int or long object)
				'f' -> serialize as single (requires float object)
		"""

		self.files = []

		return self._serialize ( obj, flags )
		

	def _serialize ( self, obj, flags = None ):
		if obj is None:
			ret = "N"

		elif isinstance ( obj, bool ):
			ret = self._serialize_bool ( obj )

		elif isinstance ( obj, str ):
			ret = self._serialize_string ( obj )
			
		elif isinstance ( obj, unicode ):
			ret = self._serialize_unicode ( obj )
			
		elif isinstance ( obj, int ):
			
			if flags:
				unsigned = ( 'u' in flags )

				if 'b' in flags:
					ret = self._serialize_int8 ( obj, unsigned )
				elif 'h' in flags:
					ret = self._serialize_int16 ( obj, unsigned )
				elif 'l' in flags:
					ret = self._serialize_int64 ( obj, unsigned )
				else:
					ret = self._serialize_int32 ( obj, unsigned )
			else:
				ret = self._serialize_int32 ( obj )

		elif isinstance ( obj, long ):

			if flags:
				unsigned = ( 'u' in flags )

				if 'b' in flags:
					ret = self._serialize_int8 ( obj, unsigned )
				elif 'h' in flags:
					ret = self._serialize_int16 ( obj, unsigned )
				elif 'i' in flags:
					ret = self._serialize_int32 ( obj, unsigned )
				elif 'l' in flags:
					ret = self._serialize_int64 ( obj, unsigned )
				else:
					ret = self._serialize_bigint ( obj )

			else:
				if obj >= -2147483648 and obj <= 2147483647:
					ret = self._serialize_int32 ( obj, 0 )
				elif obj >= 0 and obj <= 4294967295:
					ret = self._serialize_int32 ( obj, 1 )
				else:
					ret = self._serialize_bigint ( obj )

		elif isinstance ( obj, datetime ):
			ret = self._serialize_datetime ( obj )

		elif isinstance ( obj, float ):
			single = flags and ( 'f' in flags )
			ret = self._serialize_float ( obj, single )

		elif isinstance ( obj, tuple ) or isinstance ( obj, list ):
			ret = self._serialize_list ( obj )

		elif isinstance ( obj, dict ):
			ret = self._serialize_dict ( obj )

		elif isinstance ( obj, file ) or isinstance ( obj, CTPTempFile ):
			ret = self._serialize_file ( obj )
			
		else:
			if self._relaxed:
				ret = self._serialize_python ( obj )
			else:
				raise SerializationError, "cannot serialize '%s' object" % type ( obj ).__name__

		return ret


	def _deserialize_string ( self, s ):
		l = len ( s.getvalue () ) - s.tell ()
		if l < 4: raise SerializationError, "decoding error"
		slen = struct.unpack ( '>I', s.read ( 4 ) ) [ 0 ]
		l -= 4
		if l < slen: raise SerializationError, "decoding error"
		return ( s.read ( slen ), 5 + slen )


	def _deserialize_unicode ( self, s ):
		l = len ( s.getvalue () ) - s.tell ()
		if l < 4: raise SerializationError, "decoding error"
		slen = struct.unpack ( '>I', s.read ( 4 ) ) [ 0 ]
		l -= 4
		if l < slen: raise SerializationError, "decoding error"
		
		bin = s.read ( slen )
		
		a = array ( 'u' )
		a.fromstring ( bin )
		if _little_endian: a.byteswap ()

		return ( a.tounicode (), 5 + slen )


	def _deserialize_int8 ( self, s, type ):
		l = len ( s.getvalue () ) - s.tell ()
		if l < 1: raise SerializationError, "decoding error"
		return ( struct.unpack ( '>' + type, s.read ( 1 ) ) [ 0 ], 2 )

	
	def _deserialize_int16 ( self, s, type ):
		l = len ( s.getvalue () ) - s.tell ()
		if l < 2: raise SerializationError, "decoding error"
		return ( struct.unpack ( '>' + type, s.read ( 2 ) ) [ 0 ], 3 )
	
	
	def _deserialize_int32 ( self, s, type ):
		l = len ( s.getvalue () ) - s.tell ()
		if l < 4: raise SerializationError, "decoding error"
		return ( struct.unpack ( '>' + type, s.read ( 4 ) ) [ 0 ], 5 )
	
	
	def _deserialize_int64 ( self, s, type ):
		l = len ( s.getvalue () ) - s.tell ()
		if l < 8: raise SerializationError, "decoding error"
		return ( struct.unpack ( '>' + ( ( type == 'L' ) and 'Q' or 'q' ), s.read ( 8 ) ) [ 0 ], 9 )


	def _deserialize_bigint ( self, s, type ):
		l = len ( s.getvalue () ) - s.tell ()
		if l < 4: raise SerializationError, "decoding error"
		slen = struct.unpack ( '>I', s.read ( 4 ) ) [ 0 ]
		l -= 4
		if l < slen: raise SerializationError, "decoding error"

		bn = StringIO ()
		bn.write ( type )
		for c in s.read ( slen ):
			bn.write ( "%02X" % ord ( c ) )

		return ( long ( bn.getvalue (), 16 ), 5 + slen )


	def _deserialize_datetime ( self, s ):
		l = len ( s.getvalue () ) - s.tell ()
		if l < 9: raise SerializationError, "decoding error"

		vals = list ( struct.unpack ( ">HBBBBBH", s.read ( 9 ) ) )
		vals [ 6 ] *= 1000

		return ( datetime ( *vals ), 10 )


	def _deserialize_float ( self, s, type ):
		if type == 'f':
			size = 4
		else:
			size = 8
			type = 'd'
		
		l = len ( s.getvalue () ) - s.tell ()
		if l < size: raise SerializationError, "decoding error"
		return ( struct.unpack ( '>' + type, s.read ( size ) ) [ 0 ], 1 + size )


	def _deserialize_list ( self, s ):
		l = len ( s.getvalue () ) - s.tell ()
		if l < 4: raise SerializationError, "decoding error"
		count = struct.unpack ( '>I', s.read ( 4 ) ) [ 0 ]

		ret = []
		tot_size = 5

		for i in xrange ( count ):
			item, size = self._deserialize ( s )
			ret.append ( item )
			tot_size += size

		return ( ret, tot_size )


	def _deserialize_dict ( self, s ):
		l = len ( s.getvalue () ) - s.tell ()
		if l < 4: raise SerializationError, "decoding error"
		count = struct.unpack ( '>I', s.read ( 4 ) ) [ 0 ]

		ret = {}
		tot_size = 5

		for i in xrange ( count ):
			key, size = self._deserialize ( s )
			tot_size += size
			val, size = self._deserialize ( s )
			tot_size += size

			ret [ key ] = val
			
		return ( ret, tot_size )


	def _deserialize_file ( self, s ):
		fname, size = self._deserialize_string ( s )

		tfile = CTPTempFile ( fname )

		self.files.append ( ( fname, tfile ) )

		return ( tfile, size )


	def _deserialize_python ( self, s ):
		l = len ( s.getvalue () ) - s.tell ()
		if l < 4: raise SerializationError, "decoding error"
		size = struct.unpack ( '>I', s.read ( 4 ) ) [ 0 ]

		# ANDREA
		try:
			pyobj = loads ( s.read ( size ) )

		except:
			print "_deserialize_python: %s: %s" % ( sys.exc_info()[0], sys.exc_info()[1] )
			raise

		return ( pyobj, size + 5 )


	def deserialize ( self, s ):
		s = StringIO ( s )

		self.files = []
		
		return self._deserialize ( s ) [ 0 ]


	def _deserialize ( self, s ):
		if s.tell () >= len ( s.getvalue () ): return ( None, 0 )

		type = s.read ( 1 )

		if type == 'N':
			return ( None, 1 )

		elif type == '0':
			return ( False, 1 )

		elif type == '1':
			return ( True, 1 )

		elif type == 's':
			ret = self._deserialize_string ( s )

		elif type == 'S':
			ret = self._deserialize_unicode ( s )

		elif type in ( 'b', 'B' ):
			ret = self._deserialize_int8 ( s, type )
			
		elif type in ( 'h', 'H' ):
			ret = self._deserialize_int16 ( s, type )
			
		elif type in ( 'i', 'I' ):
			ret = self._deserialize_int32 ( s, type )
			
		elif type in ( 'l', 'L' ):
			ret = self._deserialize_int64 ( s, type )

		elif type in ( '+', '-' ):
			ret = self._deserialize_bigint ( s, type )

		elif type == 'd':
			ret = self._deserialize_datetime ( s )

		elif type in ( 'f', 'F' ):
			ret = self._deserialize_float ( s, type )
		
		elif type == '[':
			ret = self._deserialize_list ( s )

		elif type == '{':
			ret = self._deserialize_dict ( s )

		elif type == '/':
			ret = self._deserialize_file ( s )

		elif type == 'X':
			ret = self._deserialize_python ( s )

		else:
			raise SerializationError, "decoding error"

		return ret


if __name__ == '__main__':
	ser = CTPRPCSerializer ()
	serialize = ser.serialize
	deserialize = ser.deserialize
	
	print "NULL"
	print "null/none:                  " + serialize ( None ) + "   ---> " + str ( deserialize ( serialize ( None ) ) )
	print
	print "STRINGS"
	print "ascii (ciao):               " + serialize ( "ciao" ).__repr__ ()
	print "                 ---> " + deserialize ( serialize ( "ciao" ) )
	print "unicode (ciao \\u2641)       " + serialize ( u"ciao \u2641" ).__repr__ ()
	print "                 ---> " + deserialize ( serialize ( u"ciao \u2641" ) ).__repr__ ()
	print
	print "NUMBERS"
	print "int8 (-7):                  " + serialize ( -7, 'b' ).__repr__ () + "   ---> " + str ( deserialize ( serialize ( -7, 'b' ) ) )
	print "uint8 (7):                  " + serialize ( 7, 'ub' ).__repr__ () + "   ---> " + str ( deserialize ( serialize ( 7, 'ub' ) ) )
	print "int16 (-7):                 " + serialize ( -7, 'h' ).__repr__ () + "   ---> " + str ( deserialize ( serialize ( -7, 'h' ) ) )
	print "uint16 (7):                 " + serialize ( 7, 'uh' ).__repr__ () + "   ---> " + str ( deserialize ( serialize ( 7, 'uh' ) ) )
	print "int32 (-7):                 " + serialize ( -7 ).__repr__ () + "   ---> " + str ( deserialize ( serialize ( -7 ) ) )
	print "uint32 (7):                 " + serialize ( 7, 'u' ).__repr__ () + "   ---> " + str ( deserialize ( serialize ( 7, 'u' ) ) )
	print "int64 (-7):                 " + serialize ( -7, 'l' ).__repr__ () + "   ---> " + str ( deserialize ( serialize ( -7, 'l' ) ) )
	print "uint64 (7):                 " + serialize ( 7, 'ul' ).__repr__ () + "   ---> " + str ( deserialize ( serialize ( 7, 'ul' ) ) )
	print
	print "BIGNUM"
	print "-1311768467405634345365:    " + serialize ( -1311768467405634345365 ).__repr__ ()
	print "   ---> " + str ( deserialize ( serialize ( -1311768467405634345365 ) ) )
	print "13117320546753456453451238: " + serialize ( 13117320546753456453451238 ).__repr__ ()
	print "   ---> " + str ( deserialize ( serialize ( 13117320546753456453451238 ) ) )
	print
	print "BOOLEANS"
	print "true:                       " + serialize ( True ) + "   ---> " + str ( deserialize ( serialize ( True ) ) )
	print "false:                      " + serialize ( False ) + "   ---> " + str ( deserialize ( serialize ( False ) ) )
	print
	print "DATE/TIME"
	print "1976-10-30 14:30:12.530:    " + serialize ( datetime ( 1976, 10, 30, 14, 30, 12, 530000 ) ).__repr__ ()
	print "   ---> " + str ( deserialize ( serialize ( datetime ( 1976, 10, 30, 14, 30, 12, 530000 ) ) ) )
	print
	print "FLOATS"
	print "single (4.5):               " + serialize ( 4.5, 'f' ).__repr__ () + "   ---> " + str ( deserialize ( serialize ( 4.5, 'f' ) ) )
	print "double (4.5):               " + serialize ( 4.5 ).__repr__ () + "   ---> " + str ( deserialize ( serialize ( 4.5 ) ) )
	print
	print "LISTS"
	print "[ 'ciao', 14, None, [] ]:"
	print "   " + serialize ( [ 'ciao', 14, None, [] ] ).__repr__ ()
	print "   ---> " + str ( deserialize ( serialize ( [ 'ciao', 14, None, [] ] ) ) )
	print
	print "DICTIONARIES"
	print "{ 'pippo': 'pluto', 1: 2, 'enabled': True }:"
	print "   " + serialize ( { 'pippo': 'pluto', 1: 2, 'enabled': True } ).__repr__ ()
	print "   ---> " + str ( deserialize ( serialize ( { 'pippo': 'pluto', 1: 2, 'enabled': True } ) ) )
	print
	print "FILES"
	print "CTPRPCPickle.py:            " + serialize ( file ( "CTPRPCPickle.py" ) ).__repr__ ()
	print "   files array: " + str ( ser.files )
	print "   ---> " + str ( deserialize ( serialize ( file ( "CTPRPCPickle.py" ) ) ) )
	print "   files array: " + str ( ser.files )


	print "----------START"
	big = [ { "prova": "pipo", "prova2": "a"*1024 } ] * 50000
	s = serialize ( big )
	print "----------DESER"
	big2 = deserialize ( s )
	print "----------END"
	print big == big2

